package networkTraining;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import networkTraining.algorithms.ITrainingAlgorithm;
import networkTraining.convergence.IConvergenceTest;
import networkTraining.epoch.IEpochExecutor;
import networkTraining.epoch.TestingEpoch;
import networkTraining.inspection.IPostFinalizeInspector;
import networkTraining.inspection.IPostTestingEpochInspector;
import networkTraining.inspection.IPostTrainingEpochInspector;
import networkTraining.inspection.IPreInitializeInspector;
import networkTraining.inspection.IPreTestingEpochInspector;
import networkTraining.inspection.IPreTrainingEpochInspector;
import networkTraining.inspection.ITrainingInspector;
import neuralNetwork.INeuralNetwork;

public class NetworkTrainer implements INetworkTrainer
{
	private int maxEpochCount;

	private INeuralNetwork neuralNetwork;
	private ITrainingAlgorithm trainingAlgorithm;
	private IEpochExecutor epochExecutor;

	private IDataSetProvider dataSetProvider;

	private List<IPreInitializeInspector> preTrainingInspectors;
	private List<IPostFinalizeInspector> postTrainingInspectors;
	private List<IPreTrainingEpochInspector> preTrainingEpochInspectors;
	private List<IPostTrainingEpochInspector> postTrainingEpochInspectors;

	private List<IPreTestingEpochInspector> preTestingEpochInspectors;
	private List<IPostTestingEpochInspector> postTestingEpochInspectors;
	
	private IConvergenceTest convergenceTest;

	private TestingEpoch testingEpoch;

	public NetworkTrainer()
	{
		maxEpochCount = 10000;

		preTrainingInspectors = new ArrayList<IPreInitializeInspector>();
		postTrainingInspectors = new ArrayList<IPostFinalizeInspector>();
		preTrainingEpochInspectors = new ArrayList<IPreTrainingEpochInspector>();
		postTrainingEpochInspectors = new ArrayList<IPostTrainingEpochInspector>();
		
		preTestingEpochInspectors = new ArrayList<IPreTestingEpochInspector>();
		postTestingEpochInspectors = new ArrayList<IPostTestingEpochInspector>();
		
	}

	public NetworkTrainer(INeuralNetwork neuralNetwork, IDataSetProvider dataSetProvider,
	    ITrainingAlgorithm trainingAlgorithm, IEpochExecutor epochExecutor, int maxEpochCount,
	    IConvergenceTest convergenceTest)
	{
		this.neuralNetwork = neuralNetwork;
		this.dataSetProvider = dataSetProvider;
		this.trainingAlgorithm = trainingAlgorithm;
		this.epochExecutor = epochExecutor;
		this.maxEpochCount = maxEpochCount;
		this.convergenceTest = convergenceTest;

		preTrainingInspectors = new ArrayList<IPreInitializeInspector>();
		postTrainingInspectors = new ArrayList<IPostFinalizeInspector>();
		preTrainingEpochInspectors = new ArrayList<IPreTrainingEpochInspector>();
		postTrainingEpochInspectors = new ArrayList<IPostTrainingEpochInspector>();
	}

	@Override
	public INeuralNetwork getNeuralNetwork()
	{
		return neuralNetwork;
	}

	@Override
	public void setNeuralNetwork(INeuralNetwork neuralNetwork)
	{
		this.neuralNetwork = neuralNetwork;
	}

	@Override
	public IDataSetProvider getDataSetProvider()
	{
		return dataSetProvider;
	}

	@Override
	public void setDataSetProvider(IDataSetProvider dataSetProvider)
	{
		this.dataSetProvider = dataSetProvider;
	}

	@Override
	public ITrainingAlgorithm getTrainingAlgorithm()
	{
		return trainingAlgorithm;
	}

	@Override
	public void setTrainingAlgorithm(ITrainingAlgorithm trainingAlgorithm)
	{
		this.trainingAlgorithm = trainingAlgorithm;
	}

	@Override
	public void applyTrainingAlgorithm(Map<String, Double> errorMap)
	{
		trainingAlgorithm.trainNeuralNetwork(errorMap);
	}

	@Override
	public int getMaxEpochCount()
	{
		return maxEpochCount;
	}

	@Override
	public void setMaxEpochCount(int maxEpochCount)
	{
		this.maxEpochCount = maxEpochCount;
	}

	@Override
	public IEpochExecutor geteEpochExecutor()
	{
		return epochExecutor;
	}

	@Override
	public void setEpochExecutor(IEpochExecutor epochExecutor)
	{
		this.epochExecutor = epochExecutor;
		epochExecutor.setNetworkTrainer(this);
	}

	@Override
	public IConvergenceTest getConvergenceTest()
	{
		return convergenceTest;
	}

	@Override
	public void setConvergenceTest(IConvergenceTest convergenceTest)
	{
		this.convergenceTest = convergenceTest;
	}

	@Override
	public List<IPreInitializeInspector> getPreTrainingInspectors()
	{
		return preTrainingInspectors;
	}

	@Override
	public void addTrainingInspector(ITrainingInspector trainingInspector)
	{
		if (trainingInspector instanceof IPreInitializeInspector)
			preTrainingInspectors.add((IPreInitializeInspector) trainingInspector);

		if (trainingInspector instanceof IPostFinalizeInspector)
			postTrainingInspectors.add((IPostFinalizeInspector) trainingInspector);

		if (trainingInspector instanceof IPreTrainingEpochInspector)
			preTrainingEpochInspectors.add((IPreTrainingEpochInspector) trainingInspector);

		if (trainingInspector instanceof IPostTrainingEpochInspector)
			postTrainingEpochInspectors.add((IPostTrainingEpochInspector) trainingInspector);
		
		
		if (trainingInspector instanceof IPreTestingEpochInspector)
			preTestingEpochInspectors.add((IPreTestingEpochInspector) trainingInspector);

		if (trainingInspector instanceof IPostTestingEpochInspector)
			postTestingEpochInspectors.add((IPostTestingEpochInspector) trainingInspector);
	}

	@Override
	public List<IPostFinalizeInspector> getPostTrainingInspectors()
	{
		return postTrainingInspectors;
	}

	@Override
	public List<IPreTrainingEpochInspector> getPreTrainingEpochInspectors()
	{
		return preTrainingEpochInspectors;
	}

	@Override
	public List<IPostTrainingEpochInspector> getPostTrainingEpochInspectors()
	{
		return postTrainingEpochInspectors;
	}

	@Override
	public List<IPreTestingEpochInspector> getPreTestingEpochInspectors()
	{
		return preTestingEpochInspectors;
	}

	@Override
	public List<IPostTestingEpochInspector> getPostTestingEpochInspectors()
	{
		return postTestingEpochInspectors;
	}
	
	@Override
	public void trainNetwork()
	{
		testingEpoch = new TestingEpoch();
		testingEpoch.setNetworkTrainer(this);

		for (IPreInitializeInspector preTraining : preTrainingInspectors)
			preTraining.preInitializeInspect(neuralNetwork);

		for (int i = 0; i < getMaxEpochCount(); i++)
		{
			for (IPreTrainingEpochInspector pre : preTrainingEpochInspectors)
				pre.preEpochInspect();

			epochExecutor.executeEpoch();

			for (IPostTrainingEpochInspector post : postTrainingEpochInspectors)
				post.postEpochInspect(epochExecutor);

			epochExecutor.incrementEpochIndex();

			// TODO remove cheap hack
			// ///////////////////////////////////////////////////////
			for (IPreTestingEpochInspector pre : preTestingEpochInspectors)
				pre.preEpochInspect();

			testingEpoch.executeEpoch();

			for (IPostTestingEpochInspector post : postTestingEpochInspectors)
				post.postTestingEpochInspect(testingEpoch);

			testingEpoch.incrementEpochIndex();
			// ///////////////////////////////////////////////////////

			if (convergenceTest.isNetworkConverged(epochExecutor))
				break;

		}

		for (IPostFinalizeInspector postTraining : postTrainingInspectors)
			postTraining.postFinalizeInspect();

	}

}
